import json
import os

# 假设ZhipuAI是SDK的包名，需要根据实际情况替换
from zhipuai import ZhipuAI
# 初始化API
api_key = 'YOUR_API_KEY'  # 替换为你的API密钥
client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])
# 配置模型和初始消息
model_name = 'glm-4'
initial_messages = [{"role": "system", "content": "你是一个乐于助人的助手，为用户提供有关机器学习的信息。"}]
# 多轮对话函数
def multi_round_dialogue():
    # 将初始消息添加到对话历史中
    dialogue_history = initial_messages.copy()
    while True:
        # 获取用户输入
        user_input = input("用户: ")
        if user_input.lower() == '退出':
            print("对话结束。")
            break
        # 添加用户消息到对话历史
        dialogue_history.append({"role": "user", "content": user_input})
        # 调用模型生成回复
        response = client.chat.completions.create(
            model=model_name,
            messages=dialogue_history,
            max_tokens=150  # 可以根据需要调整最大token数
        )
        # 解析模型回复
        assistant_message = response.choices[0].message
        #dialogue_history.append(assistant_message)  # 将助手消息添加到对话历史
        dialogue_history.append({"role": "assistant", "content": assistant_message.content})  # 将助手的回复添加到消息队列中
        # 打印助手回复
        print(f"助手: {assistant_message.content}")
# 开始多轮对话
multi_round_dialogue()